import rclpy
from rclpy.node import Node
import numpy as np
import open3d as o3d
from sensor_msgs.msg import PointCloud2
from sensor_msgs_py import point_cloud2
import time  # 用于计算处理时间

class RawPointCloudVisualizer(Node):
    def __init__(self):
        super().__init__('pointcloud_visualizer')
        self.subscription = self.create_subscription(
            PointCloud2,
            'simulated_pointcloud',
            self.listener_callback,
            10)
        self.subscription  # 防止未使用变量警告
        
        # 点云缩放因子，大于1放大，小于1缩小
        self.scale_factor = 1.0  # 可以根据需要调整这个值
        
        # 点的显示大小
        self.point_size = 2.0  # 可以根据需要调整这个值
        
        # 初始化Open3D可视化窗口
        self.vis = o3d.visualization.Visualizer()
        self.vis.create_window(window_name="Simulated Cable PointCloud")
        self.pcd = o3d.geometry.PointCloud()
        self.first_update = True
        
        # 设置可视化选项
        opt = self.vis.get_render_option()
        opt.background_color = [0, 0, 0]  # RGB值范围为[0,1]，黑色背景
        opt.point_size = self.point_size  # 设置点的大小

    def listener_callback(self, msg):
        """将PointCloud2消息转换为Open3D点云并显示，计算处理时间"""
        # 记录处理开始时间
        start_time = time.time()
        
        # 从ROS2消息提取点云数据（x, y, z）
        points = point_cloud2.read_points_list(msg, field_names=("x", "y", "z"), skip_nans=True)
        
        # 转换为numpy数组
        point_array = np.array(points, dtype=np.float32)
        
        # 缩放点云
        scaled_point_array = point_array * self.scale_factor
        
        # 更新Open3D点云
        self.pcd.points = o3d.utility.Vector3dVector(scaled_point_array)
        
        # 首次更新添加点云到可视化窗口
        if self.first_update:
            self.vis.add_geometry(self.pcd)
            self.first_update = False
        else:
            self.vis.update_geometry(self.pcd)
        
        # 刷新可视化
        self.vis.poll_events()
        self.vis.update_renderer()
        
        # 计算处理时间（转换为毫秒）
        processing_time = (time.time() - start_time) * 1000  # 乘以1000转换为ms
        
        # 打印点云数量、缩放因子、点大小和处理时间信息
        self.get_logger().info(
            f"可视化点云: {len(points)}个点，缩放因子: {self.scale_factor}, "
            f"点大小: {self.point_size}, 处理时间: {processing_time:.2f} ms"
        )

def main(args=None):
    rclpy.init(args=args)
    
    # 启动可视化节点
    visualizer = RawPointCloudVisualizer()
    
    try:
        rclpy.spin(visualizer)
    except KeyboardInterrupt:
        # 关闭窗口
        visualizer.vis.destroy_window()
        pass
        
    visualizer.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()
    